import torch
import torch.nn as nn
import torch.optim as optim

# 定义共享的CNN模型
class SharedCNN(nn.Module):
    def __init__(self):
        super(SharedCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, kernel_size=2)  # 输入3通道，输出8通道，卷积核大小2x2
        self.pool1 = nn.AvgPool2d(2)  # 平均池化2x2
        self.conv2 = nn.Conv2d(8, 16, kernel_size=2)  # 输入8通道，输出16通道
        self.pool2 = nn.AvgPool2d(2)  # 平均池化2x2

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))  # 第一层卷积+池化
        x = self.pool2(torch.relu(self.conv2(x)))  # 第二层卷积+池化
        return x.view(x.size(0), -1)  # 展平为一维向量

# 定义全连接网络
class FCNetwork(nn.Module):
    def __init__(self, input_dim):
        super(FCNetwork, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)  # 输入特征的维度根据CNN的输出调整
        self.fc2 = nn.Linear(128, 1)  # 输出一个标量

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)  # 输出一个值

# 定义整个模型，包括共享的CNN和5个全连接网络
class Model(nn.Module):
    def __init__(self, cnn_output_dim):
        super(Model, self).__init__()
        self.cnn = SharedCNN()
        # 动态设置全连接网络的输入维度
        self.fc_networks = nn.ModuleList([FCNetwork(cnn_output_dim) for _ in range(5)])

    def forward(self, x):
        cnn_output = self.cnn(x)  # 获取CNN输出
        fc_outputs = [fc(cnn_output) for fc in self.fc_networks]  # 各全连接网络的输出
        return torch.cat(fc_outputs, dim=1), fc_outputs  # 返回拼接后的输出和单独的全连接输出

# 打印 CNN 和每个全连接网络的梯度
def print_gradients(model):
    print("\n=== Gradients ===")
    # 打印 CNN 的梯度
    for name, param in model.cnn.named_parameters():
        if param.requires_grad and param.grad is not None:
            print(f"{name}: {param.grad.norm():.4f}")

    # 打印每个全连接网络的梯度
    for idx, fc in enumerate(model.fc_networks):
        print(f"\nFC Network {idx + 1}:")
        for name, param in fc.named_parameters():
            if param.requires_grad and param.grad is not None:
                print(f"  {name}: {param.grad.norm():.4f}")
    print("=================\n")

# 更新 CNN 的训练步骤
def train_cnn_only(model, inputs, targets, cnn_optimizer):
    model.train()

    cnn_optimizer.zero_grad()

    # 获取模型输出
    outputs, _ = model(inputs)

    # 计算总损失（只更新 CNN）
    total_loss = 0
    for i in range(outputs.size(1)):  # 遍历所有输出维度
        total_loss += nn.MSELoss()(outputs[:, i], targets[:, i])

    total_loss.backward()  # 反向传播
    cnn_optimizer.step()  # 更新 CNN 的参数

    return total_loss.item(),outputs

# 更新全连接网络的训练步骤
def train_fc_only(model, inputs, targets, fc_optimizer, update_indices, outputs):
    model.train()

    fc_optimizer.zero_grad()

    # 获取模型输出
    outputs, _ = model(inputs)

    # 计算全连接网络的损失（仅更新指定的网络）
    total_loss = 0
    fc_losses = []
    for i in range(len(model.fc_networks)):
        if i in update_indices:
            loss = nn.MSELoss()(outputs[:, i], targets[:, i])
            total_loss += loss
            fc_losses.append(loss.item())
        else:
            fc_losses.append(None)

    total_loss.backward()  # 反向传播
    fc_optimizer.step()  # 更新全连接网络的参数

    return total_loss.item(), fc_losses

# 测试 CNN 输出维度，确保与全连接层输入匹配
test_input = torch.randn(1, 3, 32, 32)  # 单张32x32 RGB图片
cnn = SharedCNN()
cnn_output = cnn(test_input)
cnn_output_dim = cnn_output.shape[1]  # 获取展平后的特征维度
print(f"CNN 输出维度: {cnn_output_dim}")  # 确保维度正确

# 初始化模型
model = Model(cnn_output_dim)

# 定义两个优化器
cnn_optimizer = optim.Adam(model.cnn.parameters(), lr=0.001)  # 优化 CNN 的参数
fc_optimizer = optim.Adam(model.fc_networks.parameters(), lr=0.001)  # 优化全连接网络的参数

# 模拟训练数据
inputs = torch.randn(10, 3, 32, 32)  # 输入: 10张32x32的RGB图片
targets = torch.tensor([[0.2, 0.5, 0.6, 0.3, 0.4] for _ in range(10)])  # 每张图片的真实标签

# 模拟训练过程
for epoch in range(2):  # 训练2轮用于测试
    print(f"\nEpoch {epoch + 1}")

    # 第一步：只更新 CNN
    cnn_loss,outputs = train_cnn_only(model, inputs, targets, cnn_optimizer)
    print(f"Total Loss (CNN only): {cnn_loss:.4f}")
    print_gradients(model)

    # 第二步：只更新第1、2个全连接网络
    update_indices = [0, 1]
    fc_loss, fc_losses = train_fc_only(model, inputs, targets, fc_optimizer, update_indices,outputs)
    print(f"Total Loss (FC 1, 2 only): {fc_loss:.4f}")

    # 打印每个全连接网络的损失
    for idx, loss in enumerate(fc_losses):
        if loss is not None:
            print(f"FC Layer {idx + 1} Loss: {loss:.4f}")
        else:
            print(f"FC Layer {idx + 1} Loss: Not Updated")
    print_gradients(model)
